from surprise import Dataset, Reader
import pandas as pd

# 加载数据
ratings = pd.read_csv('u.data', sep='\t', names=['user','item','rating','timestamp'])
reader = Reader(rating_scale=(1, 5))
data = Dataset.load_from_df(ratings[['user', 'item', 'rating']], reader)


from surprise import KNNWithMeans
from surprise import accuracy
from surprise.model_selection import train_test_split

# 划分训练集和测试集
trainset, testset = train_test_split(data, test_size=0.2)

# 使用基于用户的协同过滤
sim_options = {
    'name': 'pearson',
    'user_based': True  # 基于用户
}

algo = KNNWithMeans(k=40, sim_options=sim_options)
algo.fit(trainset)

# 预测
predictions = algo.test(testset)

# 评估
accuracy.rmse(predictions)
accuracy.mae(predictions)